package com.maschinenstuermer.profiler;

import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.lang.management.ManagementFactory;
import java.lang.reflect.Array;
import java.lang.reflect.Field;
import java.lang.reflect.Modifier;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.IdentityHashMap;
import java.util.List;
import java.util.Map;
import java.util.Stack;
import java.util.logging.Level;
import java.util.logging.Logger;

import javax.management.MBeanServer;
import javax.management.ObjectName;

import com.google.common.base.Predicate;
import com.google.common.base.Predicates;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.maschinenstuermer.profiler.config.ConfigFactory;
import com.maschinenstuermer.profiler.model.TraceModelBuilder;
import com.maschinenstuermer.profiler.timer.TimerFactory;
import com.maschinenstuermer.profiler.transform.ClassTransfomerTemplate;
import com.maschinenstuermer.profiler.transform.JdbcClassTransformer;
import com.maschinenstuermer.profiler.transform.ProfilingClassTransformer;
import com.maschinenstuermer.profiler.transform.TransactionClassTransformer;

/**
 * This class is the profiler main entry point.
 *  
 * @author km
 */
public class ProfilerAgent {
	private static final int RETRANSFORM_CHUNKS = 25;
	private static Logger logger = Logger.getLogger(ProfilerAgent.class.getName());
	
	private static final class IsModifiableClass implements Predicate<Class<?>> {
		private final InstrumentationJdk6 instrumentation;

		private IsModifiableClass(final InstrumentationJdk6 instrumentation) {
			this.instrumentation = instrumentation;
		}

		@Override
		public boolean apply(final Class<?> input) {
			return instrumentation.isModifiableClass(input);
		}
	}
	
	private static CommonInstrumentation commonInstrumentation;
	private static Thread shutdownHook;
	
	private static List<ClassTransfomerTemplate> registedClassFileTransfomers =
		new ArrayList<ClassTransfomerTemplate>();
	private static ProfilerControl profilerControl;
	
	public static void agentmain(final String options, final Instrumentation instrumentation) {
		final CommonInstrumentation commonInstrumentation = 
			InstrumentationProxyFactory.newProxy(instrumentation);
		initProfiler(options, commonInstrumentation);
		if (commonInstrumentation.isRetransformClassesSupported()) {
			logger.info("Retransform classes supported!");
			retransformClasses(commonInstrumentation);
		} else if (commonInstrumentation.isRedefineClassesSupported()) {
			logger.info("Redefine classes supported!");
		}
	}
	
	public static void premain(final String options, final Instrumentation instrumentation) {
		final CommonInstrumentation commonInstrumentation = 
			InstrumentationProxyFactory.newProxy(instrumentation);
		initProfiler(options, commonInstrumentation);
	}

	private static void retransformClasses(final CommonInstrumentation commonInstrumentation) {
		logger.info("Retransform classes is supported!");
		final List<Class<?>> allLoadedClasses = Arrays.asList(commonInstrumentation.getAllLoadedClasses());
		final Predicate<Class<?>> transformClass = Predicates.and(new IsModifiableClass(commonInstrumentation), 
				Predicates.or(registedClassFileTransfomers));
		final ArrayList<Class<?>> transformClasses = Lists.newArrayList(Iterables.filter(allLoadedClasses, transformClass));
		profilerControl.setClassesToProcess(transformClasses.size());
		try {
			logger.info("Retransform # classes " + transformClasses.size());
			if (transformClasses.size() > 0) {
				int processedClasses = 0;
				final List<List<Class<?>>> partitions = Lists.partition(transformClasses, RETRANSFORM_CHUNKS);
				for (final List<Class<?>> partition : partitions) {					
					commonInstrumentation.retransformClasses(partition.toArray(new Class<?>[] {}));
					processedClasses += partition.size();
					profilerControl.setProcessedClasses(processedClasses);
				}
			}
		} catch (UnmodifiableClassException e) {
			logger.log(Level.SEVERE, "Error while retransforming class!", e);
		}
	}

	private static void initProfiler(final String options,
			final CommonInstrumentation commonInstrumentation) {
		ProfilerAgent.commonInstrumentation = commonInstrumentation;
		
		TimerFactory.get();
		String destDir;
		String configFile = options;
		try {
			ConfigFactory.load(configFile);
			destDir = ConfigFactory.instance().getDestDir() + "/";
		} catch (Exception e) {
			logger.throwing(ProfilerAgent.class.getName(), "premain", e);
			throw new IllegalStateException("Could not load config file!", e);
		}
		
		final ClassTransfomerTemplate profilingClassTransformer = new ProfilingClassTransformer();
		addTransformer(commonInstrumentation, profilingClassTransformer);
		
		if (ConfigFactory.instance().isEnableJdbcTracing()) {
			final ClassTransfomerTemplate jdbcTracingClassTransformer = 
				new JdbcClassTransformer();
			addTransformer(commonInstrumentation, jdbcTracingClassTransformer);
			logger.info("Sql tracing enabled!");
		}

		if (ConfigFactory.instance().isEnableTransactionTracing()) {
			final ClassTransfomerTemplate transactionTracingTransfomer = new TransactionClassTransformer();
			addTransformer(commonInstrumentation, transactionTracingTransfomer);
		}
		
		TraceModelBuilder.startProfiling();

		shutdownHook = new ShutdownHook(destDir);
		Runtime.getRuntime().addShutdownHook(shutdownHook);
		
		try {
			final ObjectName objectName = getProfilerControlName();
			
			final MBeanServer platformMBeanServer = ManagementFactory.getPlatformMBeanServer();
			if (!platformMBeanServer.isRegistered(objectName)) {
				profilerControl = new ProfilerControl();
				platformMBeanServer.registerMBean(profilerControl, objectName);
			}
		} catch (Exception e) {
			logger.log(Level.SEVERE, "", e);
		} 
	}

	public static ObjectName getProfilerControlName() {
		try {
			return new ObjectName("com.maschinenstuermer:type=ProfilerControl");
		} catch (Exception e) {
			// should never happen
			return null;
		}
	}
	
	/**
	 * Required for testing.
	 * 
	 * @return
	 */
	static Thread getShutdownHook() {
		return shutdownHook;
	}

	/**
	 * Required for testing.
	 * 
	 */
	static void setInstrumentation(final Instrumentation inst) {
		commonInstrumentation = 
			InstrumentationProxyFactory.newProxy(inst);
	}
	
	public static long sizeOf(Object obj) {
		if (commonInstrumentation == null) {
			throw new IllegalStateException(
					"Instrumentation environment not initialised.");
		}
		if (isSharedFlyweight(obj)) {
			return 0;
		}
		return commonInstrumentation.getObjectSize(obj);
	}
	
	private static void addTransformer(final CommonInstrumentation commonInstrumentation,
			final ClassTransfomerTemplate classFileTransfomer) {
		if (commonInstrumentation.isRetransformClassesSupported()) {
			commonInstrumentation.addTransformer(classFileTransfomer, true);			
		} else {
			commonInstrumentation.addTransformer(classFileTransfomer);
		}
		registedClassFileTransfomers.add(classFileTransfomer);
	}

	/**
	 * Returns deep size of object, recursively iterating over its fields and
	 * superclasses.
	 */
	@SuppressWarnings("unchecked")
	public static long deepSizeOf(Object obj) {
		Map visited = new IdentityHashMap();
		Stack stack = new Stack();
		stack.push(obj);

		long result = 0;
		do {
			result += internalSizeOf(stack.pop(), stack, visited);
		} while (!stack.isEmpty());
		return result;
	}

	/**
	 * Returns true if this is a well-known shared flyweight. For example,
	 * interned Strings, Booleans and Number objects
	 */
	@SuppressWarnings("unchecked")
	private static boolean isSharedFlyweight(Object obj) {
		// optimization - all of our flyweights are Comparable
		if (obj instanceof Comparable) {
			if (obj instanceof Enum) {
				return true;
			} else if (obj instanceof String) {
				return (obj == ((String) obj).intern());
			} else if (obj instanceof Boolean) {
				return (obj == Boolean.TRUE || obj == Boolean.FALSE);
			} else if (obj instanceof Integer) {
				return (obj == Integer.valueOf((Integer) obj));
			} else if (obj instanceof Short) {
				return (obj == Short.valueOf((Short) obj));
			} else if (obj instanceof Byte) {
				return (obj == Byte.valueOf((Byte) obj));
			} else if (obj instanceof Long) {
				return (obj == Long.valueOf((Long) obj));
			} else if (obj instanceof Character) {
				return (obj == Character.valueOf((Character) obj));
			}
		}
		return false;
	}

	@SuppressWarnings("unchecked")
	private static boolean skipObject(Object obj, Map visited) {
		return obj == null || visited.containsKey(obj)
				|| isSharedFlyweight(obj);
	}

	@SuppressWarnings("unchecked")
	private static long internalSizeOf(Object obj, Stack stack, Map visited) {
		if (skipObject(obj, visited)) {
			return 0;
		}

		Class clazz = obj.getClass();
		if (clazz.isArray()) {
			addArrayElementsToStack(clazz, obj, stack);
		} else {
			// add all non-primitive fields to the stack
			while (clazz != null) {
				Field[] fields = clazz.getDeclaredFields();
				for (Field field : fields) {
					if (!Modifier.isStatic(field.getModifiers())
							&& !field.getType().isPrimitive()) {
						field.setAccessible(true);
						try {
							stack.add(field.get(obj));
						} catch (IllegalAccessException ex) {
							throw new RuntimeException(ex);
						}
					}
				}
				clazz = clazz.getSuperclass();
			}
		}
		visited.put(obj, null);
		return sizeOf(obj);
	}

	@SuppressWarnings("unchecked")
	private static void addArrayElementsToStack(Class clazz, Object obj,
			Stack stack) {
		if (!clazz.getComponentType().isPrimitive()) {
			int length = Array.getLength(obj);
			for (int i = 0; i < length; i++) {
				stack.add(Array.get(obj, i));
			}
		}
	}
}
